import numpy as np  # NumPy是一个用于科学计算的基础包,用于处理大型多维数组和矩阵
import faiss  # FAISS库用于高效的相似度搜索和稠密向量的聚类

# 定义FaissKNeighbors类，用于执行基于FAISS的K近邻搜索
class FaissKNeighbors:
    # 类初始化函数：初始化k值，FAISS资源对象res，以及用于存储数据的索引
    def __init__(self, k=1, res=None):
        self.index = None  # 用于存储训练数据的索引
        self.y = None  # 用于存储训练数据的标签
        self.k = k  # 最近邻个数
        self.res = res  # FAISS GPU资源对象

    # 训练函数：将训练数据加入到FAISS索引中
    def fit(self, X, y):
        # 初始化 self.index 为一个FAISS索引: IndexFlatL2, 该索引使用欧氏距离进行搜索
        # 如果有GPU资源对象，则将索引转移到GPU上
        # self.index = ...
        # 确保数据为 float32 - 首先将X转换为numpy数组
        X = np.array(X)
        xb = X.astype(np.float32)
        # 维度
        if xb.ndim == 1:
            # 单个样本，转换为二维
            xb = xb.reshape(1, -1)
        d = xb.shape[1]

        # 创建 CPU 索引（L2 距离）
        cpu_index = faiss.IndexFlatL2(d)

        # 如提供 GPU 资源，则尝试将索引移动到 GPU 上
        if self.res is not None:
            try:
                self.index = faiss.index_cpu_to_gpu(self.res, 0, cpu_index)
            except Exception:
                # 若转移失败则退回 CPU 索引
                self.index = cpu_index
        else:
            self.index = cpu_index

        # 将向量加入索引（如果存在）
        if xb.size > 0:
            self.index.add(xb)

        # 保存标签为 numpy 数组，便于索引切片
        self.y = np.array(y)

    # 预测函数：对新的数据集X进行分类预测
    def predict(self, X):
        # 搜索X中每个向量的k个最近邻 - 首先将X转换为numpy数组
        X = np.array(X)
        distances, indices = self.index.search(X.astype(np.float32), self.k)
        votes = self.y[indices]  # 根据索引获得最近邻的标签
        # 通过投票机制得出最终预测的标签
        predictions = np.array([np.argmax(np.bincount(vote)) for vote in votes])
        return predictions

    # 评分函数：计算预测准确率
    def score(self, X, y):
        # 预测并比较预测结果和真实标签，计算准确率
        preds = self.predict(X)
        y_true = np.array(y)
        # 保证一维形式以便比较
        preds = np.ravel(preds)
        y_true = np.ravel(y_true)
        if y_true.size == 0:
            return 0.0
        correct = (preds == y_true).sum()
        return float(correct) / float(y_true.size)
